from Crypto.PublicKey import RSA
from Crypto.Signature import pkcs1_15
from Crypto.Hash import SHA256

# 生成密钥对
key = RSA.generate(2048)
private_key = key.export_key()
public_key = key.publickey().export_key()


# 1.密钥签名 数据
def private_key_signature(private_key, data):
    hash_obj = SHA256.new(data)
    signature = pkcs1_15.new(RSA.import_key(private_key)).sign(hash_obj)
    return signature


# 2.公钥验签
def public_key_verify(public_key, data, signature):
    hash_obj = SHA256.new(data)
    try:
        pkcs1_15.new(RSA.import_key(public_key)).verify(hash_obj, signature)
        print("验签成功")
        return True
    except:
        print("验签失败")
        return False


if __name__ == "__main__":
    msg = b"hello world"
    sig = private_key_signature(private_key, msg)
    result = public_key_verify(public_key, msg, sig)
